
#include "socket_lib.h"


static Server_start server_stars = NULL;

/**
 * @brief 发送数据包，解决 TCP 粘包问题。
 *        数据包格式：[包头][包体]，包头为 4 字节，表示包体长度。
 * 
 * @param sock 套接字描述符
 * @param message 待发送的消息
 * @return 成功发送的字节数；出错返回 -1
 */
int send_packet(int sock, const char *message) {
    int message_len = strlen(message);
    int packet_len = message_len + HEADER_LENGTH;

    // 分配内存用于存储完整数据包
    char *packet = malloc(packet_len);
    if (packet == NULL) {
        LOG_ERROR("Memory allocation failed");
        return -1;
    }

    // 构造数据包：将消息长度存入包头，消息内容存入包体
    *((int *)packet) = htonl(message_len); // 转为网络字节序（大端）
    memcpy(packet + HEADER_LENGTH, message, message_len);

    // 发送数据包
    int bytes_sent = send(sock, packet, packet_len, 0);
    if (bytes_sent != packet_len) {
        LOG_ERROR("Failed to send complete packet");
        free(packet);
        return -1;
    }

    free(packet);
    return bytes_sent;
}

/**
 * @brief 接收数据包，解决 TCP 粘包问题。
 *        数据包格式：[包头][包体]，包头为 4 字节，表示包体长度。
 * 
 * @param sock 套接字描述符
 * @param buffer 用于存储接收的包体内容
 * @return 成功接收的字节数；连接关闭返回 0；出错返回 -1
 */
int recv_packet(int sock, char *buffer) {
    int packet_length = 0;
    char packet_header[HEADER_LENGTH];

    // 接收包头（固定长度）
    ssize_t bytes_received = recv(sock, packet_header, HEADER_LENGTH, 0);
    if (bytes_received < 0) {
        LOG_ERROR("Failed to read packet header");
        return -1;
    } else if (bytes_received == 0) {
        LOG_DEBUG("Connection closed by peer\n");
        return 0; // 连接关闭
    } else if (bytes_received != HEADER_LENGTH) {
        LOG_ERROR("Incomplete packet header received\n");
        return -1;
    }

    // 解析包体长度（从网络字节序转为主机字节序）
    packet_length = ntohl(*((int *)packet_header));
    if (packet_length <= 0 || packet_length > (MAX_BUFFER_SIZE - 1)) {
        LOG_ERROR("Invalid packet length: %d\n", packet_length);
        return -1;
    }

    // 接收包体（指定长度）
    int total_received = 0;
    while (total_received < packet_length) {
        bytes_received = recv(sock, buffer + total_received, packet_length - total_received, 0);
        if (bytes_received < 0) {
            LOG_ERROR("Failed to read packet body");
            return -1;
        } else if (bytes_received == 0) {
            LOG_ERROR("Connection closed during packet reception\n");
            return -1;
        }

        total_received += bytes_received;
    }

    buffer[total_received] = '\0'; // 确保接收到的字符串以 '\0' 结尾
    return total_received;
}

/**
 * @brief 处理单个客户端的线程函数
 * 
 * 该函数负责处理客户端的连接，包括接收数据、处理请求、发送响应以及关闭连接。
 * 
 * @param arg SockInfo 结构体指针，包含客户端的 socket 和地址信息。
 * @return void 返回 0
 */
void client_thread(void *arg) {
    struct SockInfo *info = (struct SockInfo *)arg;
    char buf[MAX_BUFFER_SIZE];       // 用于存储接收到的数据的缓冲区
    char response[MAX_BUFFER_SIZE];  // 用于存储响应数据的缓冲区

    while (1) { 
        // 每次循环清空缓冲区，避免数据残留
        memset(buf, 0, MAX_BUFFER_SIZE);
        memset(response, 0, MAX_BUFFER_SIZE);

        // 接收客户端发送的数据
        int ret = recv_packet(info->fd, buf);
        if (ret == 0) {
            LOG_DEBUG("Client (%s:%d) disconnected! \n", 
                    inet_ntoa(info->addr.sin_addr), 
                    ntohs(info->addr.sin_port));
            break; // 客户端主动关闭连接
        } else if (ret == -1) {
            LOG_ERROR("Failed to receive data");
            break; // 数据读取失败
        }

        // 调用业务逻辑处理函数
        info->process_request(buf, response);

        // 发送响应数据到客户端
        send_packet(info->fd, response);
    }

    // 关闭客户端 socket 并释放资源
    close(info->fd);
    free(info);
}

/**
 * @brief 初始化服务器配置
 * 
 * 设置服务器的基本信息，包括 IP、端口以及最大连接数。
 * @param server SocketServer 结构体指针，用于存储服务器的配置。
 * @param host 服务器的 IP 地址。
 * @param port 服务器的监听端口。
 * @param max_connections 最大连接数。
 */
void socket_server_init(SocketServer *server, char *host, int port, int max_connections, int mode) {
    server->host = host;
    server->port = port;
    server->mode = mode;
    server->max_connections = max_connections;
    server->server_socket = -1;

    switch (mode)
    {
    case THREAD_POOL_TCP:
        server_stars = socket_server_start_threadpool;
        LOG_DEBUG("socket server start with threadpool!\n");
        break;
    case SELECT_TCP:
        server_stars = socket_server_start_select;
        LOG_DEBUG("socket server start with select!\n");
        if(server->max_connections > 1024 ){
            server->max_connections = 1024;
            LOG_ERROR("When select is used, the maximum connection cannot exceed 1024!\n");
        }

        break;
    case EPOLL_TCP:
        server_stars = socket_server_start_epoll;
        LOG_DEBUG("ssocket server start with epoll!\n");
        break; 
    default:
        break;
    }

}
void socket_server_start(SocketServer *server, client_request request_handler)
{
    if(server == NULL || request_handler == NULL)
    {
        LOG_DEBUG("server or request_handler is NULL\n");
        return;
    }
    server_stars(server, request_handler);
}
/**
 * @brief 启动服务器并处理客户端连接
 * 
 * 创建 socket，绑定地址，监听端口，并使用线程池处理客户端连接。
 * @param server SocketServer 结构体指针，包含服务器的配置。
 * @param client_handler 客户端处理函数指针，用于处理单个客户端的业务逻辑。
 */
void socket_server_start_threadpool(SocketServer *server, client_request request_handler) {
    struct sockaddr_in server_addr, client_addr;
    socklen_t client_addr_len = sizeof(client_addr);
    int client_sock;

    // 创建服务器 socket
    server->server_socket = socket(AF_INET, SOCK_STREAM, 0);
    if (server->server_socket < 0) {
        LOG_ERROR("Socket creation failed");
        exit(EXIT_FAILURE);
    }

    // 设置服务器地址和端口
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = inet_addr(server->host);
    server_addr.sin_port = htons(server->port);

    // 绑定 socket 到指定地址和端口
    if (bind(server->server_socket, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
        LOG_ERROR("Bind failed");
        close(server->server_socket);
        exit(EXIT_FAILURE);
    }

    // 开始监听客户端连接
    if (listen(server->server_socket, server->max_connections) < 0) {
        LOG_ERROR("Listen failed");
        close(server->server_socket);
        exit(EXIT_FAILURE);
    }

    LOG_DEBUG("Server is listening on %s:%d\n", server->host, server->port);

    // 创建线程池
    ThreadPool *pool = thread_pool_create(server->max_connections, 128);
    if (!pool) {
        LOG_ERROR( "Thread pool creation failed\n");
        close(server->server_socket);
        exit(EXIT_FAILURE);
    }

    LOG_DEBUG("Thread pool started, waiting for client connections...\n");

    while (1) {
        // 接受客户端连接
        client_sock = accept(server->server_socket, (struct sockaddr *)&client_addr, &client_addr_len);
        if (client_sock < 0) {
            LOG_ERROR("Accept failed");
            continue; // 跳过失败的连接，继续等待下一个
        }


        LOG_DEBUG("----------- New client connected (%s:%d) -----------\n",
           inet_ntoa(client_addr.sin_addr),
           ntohs(client_addr.sin_port));
        // 分配客户端信息
        struct SockInfo *info = malloc(sizeof(struct SockInfo));
        if (!info) {
            LOG_ERROR("Memory allocation failed");
            close(client_sock);
            continue;
        }

        info->fd = client_sock;
        info->addr = client_addr;
        info->process_request = request_handler;

        // 将任务添加到线程池
        if (thread_pool_add_task(pool, client_thread, info) != 0) {
            LOG_ERROR( "Failed to add task to thread pool\n");
            free(info); // 任务提交失败时释放内存
            close(client_sock);
        }
    }

    // 清理资源
    thread_pool_destroy(pool);
}

/**
 * @brief 启动服务器并处理客户端连接
 * 
 * 该函数负责启动 TCP 服务器，使用 epoll 高效管理客户端连接，
 * 并调用用户自定义的处理函数对客户端请求进行响应。
 * 
 * 主要功能：
 * 1. 创建和配置服务器套接字。
 * 2. 使用 epoll 管理多客户端的连接。
 * 3. 使用用户定义的业务逻辑处理客户端请求。
 * 
 * @param server         SocketServer 结构体指针，包含服务器的配置。
 * @param request_handler 客户端请求处理函数指针，负责处理单个客户端的业务逻辑。
 */
void socket_server_start_epoll(SocketServer *server, client_request request_handler) {
    struct sockaddr_in server_addr, client_addr;
    socklen_t client_addr_len = sizeof(client_addr);
    int client_sock;

    // 定义缓冲区，用于处理客户端数据
    char buf[MAX_BUFFER_SIZE];
    char response[MAX_BUFFER_SIZE];

    // 创建服务器 socket
    server->server_socket = socket(AF_INET, SOCK_STREAM, 0);
    if (server->server_socket < 0) {
        LOG_ERROR("Socket creation failed");
        exit(EXIT_FAILURE);
    }

    // 配置服务器地址
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = inet_addr(server->host); // 将主机地址转为网络字节序
    server_addr.sin_port = htons(server->port);            // 将端口号转为网络字节序

    // 绑定 socket 到指定地址和端口
    if (bind(server->server_socket, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
        LOG_ERROR("Bind failed");
        close(server->server_socket);
        exit(EXIT_FAILURE);
    }

    // 开始监听客户端连接
    if (listen(server->server_socket, server->max_connections) < 0) {
        LOG_ERROR("Listen failed");
        close(server->server_socket);
        exit(EXIT_FAILURE);
    }

    LOG_DEBUG("Server is listening on %s:%d\n", server->host, server->port);

    // 创建 epoll 实例
    int epoll_fd = epoll_create1(0);
    if (epoll_fd == -1) {
        LOG_ERROR("epoll_create1 failed");
        close(server->server_socket);
        exit(EXIT_FAILURE);
    }

    // 分配内存用于存储 epoll 事件
    struct epoll_event event;
    struct epoll_event *events = (struct epoll_event *)malloc(sizeof(struct epoll_event) * server->max_connections);

    // 将服务器套接字添加到 epoll 实例中
    event.events = EPOLLIN; // 监听可读事件
    event.data.fd = server->server_socket;
    if (epoll_ctl(epoll_fd, EPOLL_CTL_ADD, server->server_socket, &event) == -1) {
        LOG_ERROR("epoll_ctl failed");
        close(server->server_socket);
        close(epoll_fd);
        exit(EXIT_FAILURE);
    }



    // 主循环：等待和处理事件
    while (1) {
        // 等待 epoll 事件触发
        int event_count = epoll_wait(epoll_fd, events, server->max_connections, -1);
        if (event_count == -1) {
            LOG_ERROR("epoll_wait failed");
            continue;
        }

        // 遍历所有触发的事件
        for (int i = 0; i < event_count; i++) {
            int event_fd = events[i].data.fd;

            // 如果是服务器套接字，处理新的客户端连接
            if (event_fd == server->server_socket) {
                client_sock = accept(server->server_socket, (struct sockaddr *)&client_addr, &client_addr_len);
                if (client_sock < 0) {
                    LOG_ERROR("Accept failed");
                    continue;
                }

                LOG_DEBUG("----------- New client connected (%s:%d) -----------\n",
                inet_ntoa(client_addr.sin_addr),
                ntohs(client_addr.sin_port));

                // 添加新连接到 epoll 事件
                event.events = EPOLLIN; // 监听可读事件
                event.data.fd = client_sock;
                if (epoll_ctl(epoll_fd, EPOLL_CTL_ADD, client_sock, &event) == -1) {
                    LOG_ERROR("epoll_ctl add client failed");
                    close(client_sock);
                }
            } 
            // 如果是客户端套接字，处理数据传输
            else if (events[i].events & EPOLLIN) {
                // 清空缓冲区，准备处理数据
                memset(buf, 0, MAX_BUFFER_SIZE);
                // 接收数据
                int ret = recv_packet(event_fd, buf);
                getpeername(event_fd, (struct sockaddr *)&client_addr, &client_addr_len);
                if (ret == 0) {
                    // 客户端断开连接
                    LOG_DEBUG("Client (%s:%d) disconnected!\n",
                           inet_ntoa(client_addr.sin_addr), ntohs(client_addr.sin_port));
                    close(event_fd);
                    epoll_ctl(epoll_fd, EPOLL_CTL_DEL, event_fd, NULL);
                    continue;
                } else if (ret == -1) {
                    LOG_ERROR("Failed to receive data");
                    continue;
                }

                // 清空缓冲区，准备处理数据
                memset(response, 0, MAX_BUFFER_SIZE);

                // 调用用户自定义的请求处理逻辑
                request_handler(buf, response);

                // 发送响应数据到客户端
                send_packet(event_fd, response);
            } 
            // 处理错误或挂起事件
            else if (events[i].events & (EPOLLHUP | EPOLLERR)) {
                LOG_DEBUG("Error or hang up on socket %d\n", event_fd);
                close(event_fd);
                epoll_ctl(epoll_fd, EPOLL_CTL_DEL, event_fd, NULL);
            }
        }
    }

    // 释放资源并关闭文件描述符
    free(events);
    close(epoll_fd);
}

/**
 * @brief 启动服务器并处理客户端连接（基于 select 实现）
 * 
 * 该函数启动 TCP 服务器，使用 select 实现 I/O 多路复用，
 * 并调用用户自定义的处理函数对客户端请求进行响应。
 * 
 * @param server         SocketServer 结构体指针，包含服务器的配置。
 * @param request_handler 客户端请求处理函数指针，负责处理单个客户端的业务逻辑。
 */
void socket_server_start_select(SocketServer *server, client_request request_handler) {
    struct sockaddr_in server_addr, client_addr;
    socklen_t client_addr_len = sizeof(client_addr);
    int client_socket;

    // 定义缓冲区，用于处理客户端数据
    char buf[MAX_BUFFER_SIZE];
    char response[MAX_BUFFER_SIZE];

    // 创建服务器 socket
    server->server_socket = socket(AF_INET, SOCK_STREAM, 0);
    if (server->server_socket < 0) {
        LOG_ERROR("Socket creation failed");
        exit(EXIT_FAILURE);
    }

    // 配置服务器地址
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = inet_addr(server->host); // 将主机地址转为网络字节序
    server_addr.sin_port = htons(server->port);            // 将端口号转为网络字节序

    // 绑定 socket 到指定地址和端口
    if (bind(server->server_socket, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
        LOG_ERROR("Bind failed");
        close(server->server_socket);
        exit(EXIT_FAILURE);
    }

    // 开始监听客户端连接
    if (listen(server->server_socket, server->max_connections) < 0) {
        LOG_ERROR("Listen failed");
        close(server->server_socket);
        exit(EXIT_FAILURE);
    }

    LOG_DEBUG("Server is listening on %s:%d\n", server->host, server->port);

    // 初始化 fd 集合
    fd_set master_set, read_set;
    int max_fd;
    FD_ZERO(&master_set);
    FD_SET(server->server_socket, &master_set);
    max_fd = server->server_socket;

    // 主循环：等待和处理事件
    while (1) {
        read_set = master_set; // 复制 master_set 到 read_set

        // 使用 select 监控文件描述符集合
        if (select(max_fd + 1, &read_set, NULL, NULL, NULL) < 0) {
            LOG_ERROR("Select failed");
            break;
        }

        // 遍历所有可能的文件描述符
        for (int fd = 0; fd <= max_fd; fd++) {
            if (FD_ISSET(fd, &read_set)) {
                // 处理新连接
                if (fd == server->server_socket) {
                    client_socket = accept(server->server_socket, (struct sockaddr *)&client_addr, &client_addr_len);
                    if (client_socket < 0) {
                        LOG_ERROR("Accept failed");
                        continue;
                    }

                    LOG_DEBUG("----------- New client connected (%s:%d) -----------\n",
                    inet_ntoa(client_addr.sin_addr),
                    ntohs(client_addr.sin_port));

                    FD_SET(client_socket, &master_set); // 添加到 master_set
                    if (client_socket > max_fd) {
                        max_fd = client_socket; // 更新最大文件描述符
                    }
                } 
                // 处理客户端数据
                else {
                    // 清空缓冲区，准备处理数据
                    memset(buf, 0, MAX_BUFFER_SIZE);
                    // 接收数据
                    int ret = recv_packet(fd, buf);
                    getpeername(fd, (struct sockaddr *)&client_addr, &client_addr_len);
                    if (ret == 0) {
                        // 客户端断开连接
                        LOG_DEBUG("Client (%s:%d) disconnected!\n",
                            inet_ntoa(client_addr.sin_addr), ntohs(client_addr.sin_port));
                        close(fd);
                        FD_CLR(fd, &master_set); // 从集合中移除
                        continue;
                    } else if (ret == -1) {
                        LOG_ERROR("Failed to receive data");
                        continue;
                    }

                    // 清空缓冲区，准备处理数据
                    memset(response, 0, MAX_BUFFER_SIZE);

                    // 调用用户自定义的请求处理逻辑
                    request_handler(buf, response);

                    // 发送响应数据到客户端
                    send_packet(fd, response);

                }
            }
        }
    }

}

/**
 * @brief 停止服务器
 * 
 * 关闭服务器 socket。
 * @param server SocketServer 结构体指针。
 */
void socket_server_stop(SocketServer *server) {
    
    if (server->server_socket != -1) {
        close(server->server_socket);
        LOG_DEBUG("Server closed\n");
    }
}

/**
 * @brief 初始化 SocketClient 结构体。
 * 
 * @param client SocketClient 实例
 * @param host 服务器地址
 * @param port 服务器端口
 */
void socket_client_init(SocketClient *client, char *host, int port) {
    if (!client || !host) {
        LOG_ERROR("Client or host pointer is NULL\n");
        return;
    }

    client->host = host;
    client->port = port;
    client->client_socket = -1;
}

/**
 * @brief 连接到服务器。
 * 
 * @param client SocketClient 实例
 */
void socket_client_connect(SocketClient *client) {
    if (!client) {
        LOG_ERROR("Client pointer is NULL\n");
        return;
    }

    struct sockaddr_in server_addr = {0};

    client->client_socket = socket(AF_INET, SOCK_STREAM, 0);
    if (client->client_socket < 0) {
        LOG_ERROR("Socket creation failed");
        exit(EXIT_FAILURE);
    }

    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = inet_addr(client->host);
    server_addr.sin_port = htons(client->port);

    if (connect(client->client_socket, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
        LOG_ERROR("Connection failed");
        close(client->client_socket);
        exit(EXIT_FAILURE);
    }

    LOG_DEBUG("Connected to server %s:%d\n", client->host, client->port);
}

/**
 * @brief 发送消息到服务器并接收响应。
 * 
 * @param client SocketClient 实例
 * @param message 待发送的消息
 * @param response 用于存储服务器响应的指针
 */
void socket_client_send(SocketClient *client, char *message, char **response) {
    if (!client || !message || !response) {
        LOG_ERROR("Client, message, or response pointer is NULL\n");
        return;
    }

    if (send_packet(client->client_socket, message) < 0) {
        LOG_DEBUG("Failed to send message\n");
        return;
    }

    // 分配内存用于接收服务器响应
    char *buffer = malloc(MAX_BUFFER_SIZE);
    if (!buffer) {
        LOG_ERROR("Memory allocation failed");
        return;
    }

    int packet_size = recv_packet(client->client_socket, buffer);
    if (packet_size > 0) {
        *response = buffer;
    } else {
        free(buffer);
        *response = NULL;
    }
}

/**
 * @brief 关闭与服务器的连接。
 * 
 * @param client SocketClient 实例
 */
void socket_client_close(SocketClient *client) {
    if (!client) {
        LOG_ERROR("Client pointer is NULL\n");
        return;
    }

    if (client->client_socket != -1) {
        close(client->client_socket);
        client->client_socket = -1;
    }
}
